import os
import cv2
import torch
import numpy as np
import pandas as pd
import torch.nn as nn
import torchvision
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, models
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
import time
import datetime

# 设置中文字体显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]

# 数据层：OCT数据集处理
class OCTDataset(Dataset):
    def __init__(self, img_paths, labels, transform=None):
        """
        初始化OCT数据集
        :param img_paths: 图像路径列表
        :param labels: 标签列表
        :param transform: 数据变换
        """
        self.img_paths = img_paths
        self.labels = labels
        self.transform = transform
        # 类别名称映射（0=正常，1=黄斑变性，2=糖尿病视网膜病变，3=青光眼，4=视网膜脱离）
        self.class_names = ["正常", "黄斑变性", "糖尿病视网膜病变", "青光眼", "视网膜脱离"]
        
    def __len__(self):
        return len(self.img_paths)
    
    def __getitem__(self, idx):
        # 读取图像
        img_path = self.img_paths[idx]
        img = cv2.imread(img_path)
        
        # 确保图像读取成功
        if img is None:
            raise ValueError(f"无法读取图像: {img_path}")
            
        # 转换为RGB格式（OCT有时是灰度图）
        if len(img.shape) == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
        else:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
        # 调整大小（降低分辨率以适应基层设备）
        img = cv2.resize(img, (256, 256))
        
        # 应用数据变换
        if self.transform:
            img = self.transform(img)
            
        # 获取标签
        label = self.labels[idx]
        
        return img, label, img_path

# 数据增强和预处理
def get_transforms():
    """创建训练和验证的数据变换"""
    train_transform = transforms.Compose([
        transforms.ToPILImage(),
        transforms.RandomHorizontalFlip(),  # 随机水平翻转
        transforms.RandomVerticalFlip(),    # 随机垂直翻转
        transforms.RandomRotation(15),      # 随机旋转
        transforms.ColorJitter(brightness=0.2, contrast=0.2),  # 亮度和对比度调整
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    val_transform = transforms.Compose([
        transforms.ToPILImage(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    return train_transform, val_transform

# 模型层：轻量级OCT分类模型
class LightweightOCTClassifier(nn.Module):
    def __init__(self, num_classes=5):
        super(LightweightOCTClassifier, self).__init__()
        # 使用MobileNetV2作为基础模型，它比V3更轻量，适合基层设备
        self.base_model = models.mobilenet_v2(pretrained=True)
        
        # 冻结部分底层参数，减少训练计算量
        for param in list(self.base_model.parameters())[:-10]:
            param.requires_grad = False
            
        # 替换分类层以适应我们的类别数
        in_features = self.base_model.classifier[1].in_features
        self.base_model.classifier[1] = nn.Linear(in_features, num_classes)
        
    def forward(self, x):
        return self.base_model(x)

# 训练和评估函数
def train_model(model, train_loader, val_loader, criterion, optimizer, device, num_epochs=5):
    """
    训练模型并返回训练好的模型和训练记录
    """
    # 记录训练过程
    history = {
        'train_loss': [],
        'train_acc': [],
        'val_loss': [],
        'val_acc': []
    }
    
    best_val_acc = 0.0
    best_model_weights = None
    
    for epoch in range(num_epochs):
        print(f'Epoch {epoch+1}/{num_epochs}')
        print('-' * 10)
        
        # 训练阶段
        model.train()
        running_loss = 0.0
        running_corrects = 0
        
        for inputs, labels, _ in train_loader:
            inputs = inputs.to(device)
            labels = labels.to(device)
            
            # 清零梯度
            optimizer.zero_grad()
            
            # 前向传播
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            loss = criterion(outputs, labels)
            
            # 反向传播和优化
            loss.backward()
            optimizer.step()
            
            # 统计
            running_loss += loss.item() * inputs.size(0)
            running_corrects += torch.sum(preds == labels.data)
        
        # 计算训练集的平均损失和准确率
        epoch_loss = running_loss / len(train_loader.dataset)
        epoch_acc = running_corrects.double() / len(train_loader.dataset)
        
        history['train_loss'].append(epoch_loss)
        history['train_acc'].append(epoch_acc.item())
        
        print(f'Train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')
        
        # 验证阶段
        model.eval()
        val_running_loss = 0.0
        val_running_corrects = 0
        all_preds = []
        all_labels = []
        
        with torch.no_grad():
            for inputs, labels, _ in val_loader:
                inputs = inputs.to(device)
                labels = labels.to(device)
                
                outputs = model(inputs)
                _, preds = torch.max(outputs, 1)
                loss = criterion(outputs, labels)
                
                val_running_loss += loss.item() * inputs.size(0)
                val_running_corrects += torch.sum(preds == labels.data)
                
                # 收集所有预测和标签用于详细评估
                all_preds.extend(preds.cpu().numpy())
                all_labels.extend(labels.cpu().numpy())
        
        # 计算验证集的平均损失和准确率
        val_epoch_loss = val_running_loss / len(val_loader.dataset)
        val_epoch_acc = val_running_corrects.double() / len(val_loader.dataset)
        
        history['val_loss'].append(val_epoch_loss)
        history['val_acc'].append(val_epoch_acc.item())
        
        print(f'Val Loss: {val_epoch_loss:.4f} Acc: {val_epoch_acc:.4f}')
        print()
        
        # 保存最佳模型
        if val_epoch_acc > best_val_acc:
            best_val_acc = val_epoch_acc
            best_model_weights = model.state_dict()
    
    # 加载最佳模型权重
    model.load_state_dict(best_model_weights)
    
    # 打印详细分类报告
    class_names = ["正常", "黄斑变性", "糖尿病视网膜病变", "青光眼", "视网膜脱离"]
    print("详细分类报告:")
    print(classification_report(all_labels, all_preds, target_names=class_names))
    
    return model, history

# 应用层：医生使用的图形界面
class OCTDiagnosticApp:
    def __init__(self, root, model, class_names, device):
        self.root = root
        self.root.title("乡村卫生院眼部OCT影像辅助诊断系统")
        self.root.geometry("800x600")
        self.root.resizable(True, True)
        
        # 模型和设备
        self.model = model
        self.class_names = class_names
        self.device = device
        
        # 历史记录
        self.history = []
        
        # 数据变换
        self.transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        
        # 创建界面组件
        self._create_widgets()
        
    def _create_widgets(self):
        # 创建主框架
        main_frame = tk.Frame(self.root, padx=10, pady=10)
        main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 左侧：影像显示区
        left_frame = tk.Frame(main_frame, width=400, height=500)
        left_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=(0, 10))
        
        # 影像上传按钮
        upload_btn = tk.Button(left_frame, text="上传OCT影像", command=self.upload_image,
                              font=("SimHei", 12), height=2)
        upload_btn.pack(pady=(0, 10), fill=tk.X)
        
        # 影像显示区域
        self.image_frame = tk.LabelFrame(left_frame, text="OCT影像", padx=5, pady=5)
        self.image_frame.pack(fill=tk.BOTH, expand=True)
        
        self.image_label = tk.Label(self.image_frame, text="请上传OCT影像", font=("SimHei", 12))
        self.image_label.pack(expand=True)
        
        # 右侧：结果和操作区
        right_frame = tk.Frame(main_frame, width=400, height=500)
        right_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True)
        
        # 诊断按钮
        diagnose_btn = tk.Button(right_frame, text="开始诊断", command=self.run_diagnosis,
                                font=("SimHei", 12), height=2, bg="#4CAF50", fg="white")
        diagnose_btn.pack(pady=(0, 10), fill=tk.X)
        
        # 结果显示区域
        result_frame = tk.LabelFrame(right_frame, text="诊断结果", padx=5, pady=5)
        result_frame.pack(fill=tk.BOTH, expand=True, pady=(0, 10))
        
        self.result_text = tk.Text(result_frame, font=("SimHei", 10), wrap=tk.WORD, state=tk.DISABLED)
        self.result_text.pack(fill=tk.BOTH, expand=True)
        
        # 历史记录区域
        history_frame = tk.LabelFrame(right_frame, text="最近诊断记录", padx=5, pady=5)
        history_frame.pack(fill=tk.BOTH, expand=True)
        
        self.history_listbox = tk.Listbox(history_frame, font=("SimHei", 10))
        self.history_listbox.pack(fill=tk.BOTH, expand=True, side=tk.LEFT)
        
        scrollbar = tk.Scrollbar(history_frame, orient=tk.VERTICAL, command=self.history_listbox.yview)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        self.history_listbox.config(yscrollcommand=scrollbar.set)
        
        # 绑定双击查看历史记录
        self.history_listbox.bind("<Double-1>", self.view_history)
        
    def upload_image(self):
        """上传OCT影像"""
        file_path = filedialog.askopenfilename(
            filetypes=[("图像文件", "*.png;*.jpg;*.jpeg;*.bmp")]
        )
        
        if file_path:
            self.current_image_path = file_path
            # 显示图像
            img = Image.open(file_path)
            # 调整图像大小以适应显示区域
            img.thumbnail((350, 350))
            photo = ImageTk.PhotoImage(img)
            
            self.image_label.config(image=photo)
            self.image_label.image = photo  # 保持引用
            
            # 清空之前的结果
            self.result_text.config(state=tk.NORMAL)
            self.result_text.delete(1.0, tk.END)
            self.result_text.config(state=tk.DISABLED)
    
    def run_diagnosis(self):
        """运行诊断"""
        if not hasattr(self, 'current_image_path'):
            messagebox.showwarning("警告", "请先上传OCT影像")
            return
            
        try:
            # 显示处理中
            self.result_text.config(state=tk.NORMAL)
            self.result_text.delete(1.0, tk.END)
            self.result_text.insert(tk.END, "正在分析影像，请稍候...\n")
            self.result_text.config(state=tk.DISABLED)
            self.root.update()
            
            # 记录开始时间
            start_time = time.time()
            
            # 读取和预处理图像
            img = cv2.imread(self.current_image_path)
            
            if len(img.shape) == 2:
                img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
            else:
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                
            img = cv2.resize(img, (256, 256))
            img_tensor = self.transform(img).unsqueeze(0).to(self.device)
            
            # 模型预测
            self.model.eval()
            with torch.no_grad():
                outputs = self.model(img_tensor)
                probabilities = torch.nn.functional.softmax(outputs, dim=1)
                _, predicted = torch.max(outputs, 1)
            
            # 计算耗时
            end_time = time.time()
            elapsed_time = end_time - start_time
            
            # 获取结果
            class_idx = predicted.item()
            class_name = self.class_names[class_idx]
            confidence = probabilities[0][class_idx].item() * 100
            
            # 生成诊断建议
            diagnosis_result, risk_level = self._generate_diagnosis(class_idx, confidence)
            
            # 显示结果
            self.result_text.config(state=tk.NORMAL)
            self.result_text.delete(1.0, tk.END)
            self.result_text.insert(tk.END, f"诊断完成 (耗时: {elapsed_time:.2f}秒)\n\n")
            self.result_text.insert(tk.END, f"病变类型: {class_name}\n")
            self.result_text.insert(tk.END, f"置信度: {confidence:.2f}%\n\n")
            self.result_text.insert(tk.END, f"风险等级: {risk_level}\n\n")
            self.result_text.insert(tk.END, f"诊断建议:\n{diagnosis_result}\n")
            self.result_text.config(state=tk.DISABLED)
            
            # 保存到历史记录
            self._save_to_history(class_name, confidence, risk_level)
            
        except Exception as e:
            messagebox.showerror("错误", f"诊断过程中发生错误: {str(e)}")
    
    def _generate_diagnosis(self, class_idx, confidence):
        """根据分类结果生成诊断建议"""
        # 正常情况
        if class_idx == 0:
            if confidence > 90:
                return (
                    "未检测到明显眼部病变，建议保持良好的用眼习惯，每年进行一次常规眼部检查。",
                    "低风险"
                )
            else:
                return (
                    "未明确检测到眼部病变，但结果不确定性较高，建议3个月后复查。",
                    "中低风险"
                )
        
        # 黄斑变性
        elif class_idx == 1:
            if confidence > 85:
                return (
                    "疑似黄斑变性，这是一种常见的眼底疾病，可能影响中心视力。建议尽快转诊至上级医院进行进一步检查和治疗。",
                    "高风险"
                )
            else:
                return (
                    "可能存在黄斑变性迹象，建议2周内到上级医院进行详细检查确认诊断。",
                    "中高风险"
                )
        
        # 糖尿病视网膜病变
        elif class_idx == 2:
            if confidence > 85:
                return (
                    "疑似糖尿病视网膜病变，这是糖尿病的常见并发症。建议立即控制血糖并转诊至上级医院眼科进行详细检查和治疗。",
                    "高风险"
                )
            else:
                return (
                    "可能存在糖尿病视网膜病变迹象，建议严格控制血糖并在1周内到上级医院检查。",
                    "中高风险"
                )
        
        # 青光眼
        elif class_idx == 3:
            if confidence > 85:
                return (
                    "疑似青光眼，这是一种可能导致不可逆视力丧失的疾病。建议尽快转诊至上级医院进行眼压测量和视野检查。",
                    "高风险"
                )
            else:
                return (
                    "可能存在青光眼迹象，建议2周内到上级医院进行进一步检查确认诊断。",
                    "中高风险"
                )
        
        # 视网膜脱离
        elif class_idx == 4:
            if confidence > 80:
                return (
                    "疑似视网膜脱离，这是一种紧急情况，如不及时治疗可能导致永久性视力丧失。建议立即转诊至有条件的医院进行急诊处理。",
                    "极高风险"
                )
            else:
                return (
                    "可能存在视网膜脱离迹象，这是一种紧急情况，建议24小时内到上级医院进行详细检查。",
                    "高风险"
                )
        
        # 默认情况
        return ("无法确定诊断结果，建议立即转诊至上级医院进行详细检查。", "未知风险")
    
    def _save_to_history(self, class_name, confidence, risk_level):
        """保存诊断记录到历史"""
        timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M")
        filename = os.path.basename(self.current_image_path)
        
        # 保存完整记录
        record = {
            "timestamp": timestamp,
            "filename": filename,
            "path": self.current_image_path,
            "class_name": class_name,
            "confidence": confidence,
            "risk_level": risk_level
        }
        
        self.history.append(record)
        
        # 只保留最近10条记录
        if len(self.history) > 10:
            self.history = self.history[-10:]
        
        # 更新历史记录列表
        self.history_listbox.delete(0, tk.END)
        for i, item in enumerate(reversed(self.history)):
            display_text = f"{item['timestamp']} - {item['class_name']} - {item['filename']}"
            self.history_listbox.insert(tk.END, display_text)
    
    def view_history(self, event):
        """查看历史记录详情"""
        selected_index = self.history_listbox.curselection()
        if not selected_index:
            return
            
        # 获取选中的记录（注意列表是倒序显示的）
        index = len(self.history) - 1 - selected_index[0]
        record = self.history[index]
        
        # 显示图像
        img = Image.open(record["path"])
        img.thumbnail((350, 350))
        photo = ImageTk.PhotoImage(img)
        
        self.image_label.config(image=photo)
        self.image_label.image = photo  # 保持引用
        
        # 显示结果
        self.result_text.config(state=tk.NORMAL)
        self.result_text.delete(1.0, tk.END)
        self.result_text.insert(tk.END, f"诊断时间: {record['timestamp']}\n")
        self.result_text.insert(tk.END, f"影像文件: {record['filename']}\n\n")
        self.result_text.insert(tk.END, f"病变类型: {record['class_name']}\n")
        self.result_text.insert(tk.END, f"置信度: {record['confidence']:.2f}%\n\n")
        self.result_text.insert(tk.END, f"风险等级: {record['risk_level']}\n\n")
        
        # 重新生成诊断建议
        class_idx = self.class_names.index(record['class_name'])
        diagnosis_result, _ = self._generate_diagnosis(class_idx, record['confidence'])
        self.result_text.insert(tk.END, f"诊断建议:\n{diagnosis_result}\n")
        
        self.result_text.config(state=tk.DISABLED)

# 主函数：整合数据、模型和应用
def main():
    # 设置随机种子，保证结果可复现
    torch.manual_seed(42)
    np.random.seed(42)
    
    # 检查是否有可用的GPU，否则使用CPU（适应基层设备）
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    
    # 类别名称
    class_names = ["正常", "黄斑变性", "糖尿病视网膜病变", "青光眼", "视网膜脱离"]
    num_classes = len(class_names)
    
    # 检查是否有预训练模型，没有则训练新模型
    model_path = "oct_classifier.pth"
    
    if os.path.exists(model_path):
        # 加载预训练模型
        print("加载预训练模型...")
        model = LightweightOCTClassifier(num_classes=num_classes)
        model.load_state_dict(torch.load(model_path, map_location=device))
        model = model.to(device)
        model.eval()
    else:
        # 这里应该有数据加载和模型训练的代码
        # 实际应用中，需要替换为真实的数据集路径
        print("警告: 未找到预训练模型。在实际应用中，这里应该加载数据集并训练模型。")
        print("创建一个临时模型用于演示...")
        model = LightweightOCTClassifier(num_classes=num_classes)
        model = model.to(device)
        
        # 保存模型供下次使用
        torch.save(model.state_dict(), model_path)
    
    # 启动应用程序
    print("启动OCT影像辅助诊断系统...")
    root = tk.Tk()
    app = OCTDiagnosticApp(root, model, class_names, device)
    root.mainloop()

if __name__ == "__main__":
    main()
